﻿#include <iostream>
#include <winsock.h>
#include <fstream>
#include <time.h>
#include <string>
#include<queue>
#include "reliableUdp.h"
#pragma comment(lib, "Ws2_32.lib")
using namespace std;

char buf[200000000];

SOCKET client;

SOCKADDR_IN serverAddr, clientAddr;

bool shakeHand() {
	while (true) {
		char shake_package[2];
		shake_package[1] = SHAKE_1;
		shake_package[0] = checksum(shake_package + 1, 1);
		sendto(client, shake_package, 2, 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
		cout << "Sent Shake 1.\n";
		//第一次握手
		int begin = clock();
		char recv[2];
		int serverAddrLen = sizeof(serverAddr);
		bool lost = false;
		while (recvfrom(client, recv, 2, 0, (sockaddr*)&serverAddr, &serverAddrLen) == SOCKET_ERROR) {
			if (clock() - begin > MAX_WAIT_TIME) {
				lost = true;
				cout << "Time limit for receiving Shake 2 exceeded.\n";
				break;
			}
		}
		if (!lost && checksum(recv, 2) == 0 && recv[1] == SHAKE_2) {
			cout << "Received Shake 2.\n";
			//第三次握手
			shake_package[1] = SHAKE_3;
			shake_package[0] = checksum(shake_package + 1, 1);
			sendto(client, shake_package, 2, 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
			cout << "Sent Shake 3.\n";
			break;
		}
	}
	while (true) {
		bool hasResentShake2 = true;
		char recv[2];
		int serverAddrLen = sizeof(serverAddr);
		int begin = clock();
		while (recvfrom(client, recv, 2, 0, (sockaddr*)&serverAddr, &serverAddrLen) == SOCKET_ERROR) {
			if (clock() - begin > MAX_WAVE_TIME) {
				hasResentShake2 = false;
				break;
			}
		}
		if (hasResentShake2 && checksum(recv, 2) == 0 && recv[1] == SHAKE_2) {
			cout << "Received Shake 2.\n";
			char send[2];
			send[1] = SHAKE_3;
			send[0] = checksum(send + 1, 1);
			sendto(client, send, 2, 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
			cout << "Sent Shake 3.\n";
			continue;
		}
		break;
	}
	cout << "Shaking hand completed.\n";
	return true;
}
bool waveHand() {
	while (true) {
		char send[2];
		send[1] = WAVE_1;
		send[0] = checksum(send + 1, 1);
		sendto(client, send, 2, 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
		cout << "Sent Wave 1.\n";
		int begin = clock();
		char recv[2];
		int serverAddrLen = sizeof(serverAddr);
		bool lost = false;
		while (true) {
			int recvRtn = recvfrom(client, recv, 2, 0, (sockaddr*)&serverAddr, &serverAddrLen);
			if (recvRtn > 0 && checksum(recv, 2) == 0) {
				if (recv[1] == WAVE_2) {
					cout << "Received Wave 2.\n";
					lost = false;
					break;
				}
				if (recv[1] == WAVE_3) {
					cout<<"Received Wave 3.\n";
					lost = false;
					break;
				}
			}
			else {
				if (clock() - begin > MAX_WAIT_TIME) {
					lost = true;
					cout << "Time limit for receiving Wave 2 exceeded.\n";
					break;
				}
			}
		}
		if (!lost && checksum(recv, 2) == 0 && recv[1] == WAVE_2) {
			break;
		}
	}
	while (true) {
		char recv[2];
		int serverAddrLen = sizeof(serverAddr);
		int begin = clock();
		while (recvfrom(client, recv, 2, 0, (sockaddr*)&serverAddr, &serverAddrLen) == SOCKET_ERROR);
		if (checksum(recv, 2) != 0 || recv[1] != (char)WAVE_3) {
			continue;
		}
		cout << "Received Wave 3.\n";
		char send[2];
		send[1] = WAVE_4;
		send[0] = checksum(send + 1, 1);
		sendto(client, send, 2, 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
		cout << "Sent Wave 4.\n";
		break;
	}
	while (true) {
		bool hasResentWave3 = true;
		char recv[2];
		int serverAddrLen = sizeof(serverAddr);
		int begin = clock();
		while (recvfrom(client, recv, 2, 0, (sockaddr*)&serverAddr, &serverAddrLen) == SOCKET_ERROR) {
			if (clock() - begin > MAX_WAVE_TIME) {
				hasResentWave3 = false;
				break;
			}
		}
		if (hasResentWave3 && checksum(recv, 2) == 0 && recv[1] == WAVE_3) {
			cout << "Received Wave 3.\n";
			char send[2];
			send[1] = WAVE_4;
			send[0] = checksum(send + 1, 1);
			sendto(client, send, 2, 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
			cout << "Sent Wave 4.\n";
			continue;
		}
		break;
	}
	cout << "Waving hand completed.\n";
	return true;
}

bool rdtSendPkt(char* pkt, int len, bool last, unsigned char pktFlag) {
	if (len > MAX_LEN || (!last && len != MAX_LEN)) {
		cout << "[Error] Wrong packet splitted size.\n";
		return false;
	}
	char* pkg;
	int pkgLen;
	if (!last) {
		pkg = new char[len + 3];
		pkg[1] = MIDDLE;
		pkg[2] = pktFlag;
		for (int i = 3; i < len + 3; i++) {
			pkg[i] = pkt[i - 3];
		}
		pkg[0] = checksum(pkg + 1, len + 2);
		pkgLen = len + 3;
	}
	else {
		pkg = new char[len + 5];
		pkg[1] = LAST;
		pkg[2] = pktFlag;
		pkg[3] = len >> 8;
		pkg[4] = len & 0xFF;
		for (int i = 5; i < len + 5; i++) {
			pkg[i] = pkt[i - 5];
		}
		pkg[0] = checksum(pkg + 1, len + 4);
		pkgLen = len + 5;
	}
	cout << "checksum = " << (int)((unsigned char)pkg[0]) << ": \n";
	sendto(client, pkg, pkgLen, 0, (sockaddr*)&serverAddr, sizeof(serverAddr));
	delete[]pkg;
	return true;
}

bool rdtSender(char* file, int len) {
	int sentNum = 0;//已经发送的数量
	int sentRcvdNum = 0;//已经成功发送的数量
	int pktNum = len / MAX_LEN + (len % MAX_LEN != 0);
	bool isPKT0 = true;//控制pktFlag
	while (sentRcvdNum < pktNum) {
		if (!rdtSendPkt(file + sentRcvdNum * MAX_LEN, sentRcvdNum == pktNum - 1 ? len - (pktNum - 1) * MAX_LEN : MAX_LEN, sentRcvdNum == pktNum - 1, isPKT0 ? PKT0 : PKT1)) {
			cout << "[Error] rdtSendPkt failure.\n";
			return false;
		}
		sentNum++;
		cout << "Sent packet " << sentNum << ", " << (isPKT0 ? "PKT0" : "PTK1") <<", Len = "<< (sentRcvdNum == pktNum - 1 ? len - (pktNum - 1) * MAX_LEN : MAX_LEN) << endl;
		int begin = clock();
		//接受ACK
		char recv[2];
		int serverAddrLen = sizeof(serverAddr);
		while (true) {
			int recvRtn = recvfrom(client, recv, 2, 0, (sockaddr*)&serverAddr, &serverAddrLen);
			if (recvRtn > 0 && checksum(recv, 2) == 0) {
				if (recv[1] == (isPKT0 ? ACK0 : ACK1)) {
					cout << "Received ACK for packet " << sentNum << ", " << (isPKT0 ? "ACK0" : "ACK1") << endl;
					isPKT0 = !isPKT0;
					sentRcvdNum++;
					break;
				}
				else {
					if (recv[1] == NAK) {
						cout << "[Warning] Received NAK.\n";
						break;
					}
					else {
						cout << "[Error] Received something not ACK or NAK.\n";
						return false;
					}
				}
			}
			else {
				if (clock() - begin > MAX_WAIT_TIME) {
					cout << "[Warning] Time limit exceeded. Resend.\n";
					break;
				}
			}
		}
	}
	cout << "Sent successfully with resent number " << sentNum - sentRcvdNum << endl;
	return true;
}

int main() {
	WSADATA wsadata;
	if (WSAStartup(MAKEWORD(2, 2), &wsadata)) {
		printf("error!");
		return 0;
	}

	serverAddr.sin_family = AF_INET;
	serverAddr.sin_port = htons(8000);
	serverAddr.sin_addr.s_addr = inet_addr("127.0.0.1");

	clientAddr.sin_family = AF_INET;
	clientAddr.sin_port = htons(10005);
	clientAddr.sin_addr.s_addr = inet_addr("127.0.0.1");

	client = socket(AF_INET, SOCK_DGRAM, 0);

	int time_out = 5000;
	setsockopt(client, SOL_SOCKET, SO_RCVTIMEO, (char*)&time_out, sizeof(time_out));

	if (client == INVALID_SOCKET) {
		printf("socket of client invalid!");
		closesocket(client);
		return 0;
	}

	if (bind(client, (sockaddr*)(&clientAddr), sizeof(clientAddr)) == SOCKET_ERROR) {
		printf("bind fail");
		closesocket(client);
		WSACleanup();
		return 0;
	}
	//握手
	cout << "Trying to establish connection...\n";
	shakeHand();
	cout << "Connection with client " << inet_ntoa(serverAddr.sin_addr) << ": " << ntohs(serverAddr.sin_port) << " established.\n===================================\n";
	while (true) {
		//传文件
		string fileName;
		cout << "Please input your filename: \n";
		cin >> fileName;
		rdtSender((char*)fileName.c_str(), fileName.length());
		//exit
		if (!strcmp("/exit", fileName.c_str())) {
			break;
		}
		// 使用二进制方式 打开当前目录下的文件
		string fileDir = "C:\\Users\\16116\\Desktop\\rt\\file\\";
		string filePath = fileDir + fileName;
		ifstream in(filePath.c_str(), ifstream::binary);
		int len = 0;
		if (!in) {
			cout << "can't open the file! Please retry\n";
			return 0;
		}
		// 文件读取到buffer
		BYTE t = in.get();
		while (in) {
			buf[len++] = t;
			t = in.get();
		}
		in.close();
		rdtSender(buf, len);
	}
	waveHand();
	closesocket(client);
	cout << "Connection Aborted.\n";
	WSACleanup();
	cout << "WSA cleaning up finished.\n";
	return 0;
}